# 重新优化设计布局，左侧对话窗口保持不变，右侧不再显示历史记录，而是修改为显示对话按钮和历史记录列表。
# deepseek提示词：
# 重新设计sqlite数据库结构，采用第一个问题的时间作为唯一标识，保存第一个用户输入问题为主题，将全部对话内容作为记录，保留对话时间作为时间，对话人作为用户。
# 修改前端页面，将对话内容显示在左侧对话窗口，右侧显示历史对话列表，
# 右侧上部是对话控制区，包括4个按钮，分别是：新建按钮：用第一个问题+时间为主题，保存当前对话到数据库，清空后，创新对话；导出：将当前对话保存为csv文件；全部导出：将数据库的内容保存为csv；清空：将数据库清空
# 右侧下部是历史对话列表，将历史对话按时间逆序排列，点击历史对话列表中的对话，可以显示该对话的内容，继续开始对话。

"""
文件说明
数据库初始化：使用 init_db 函数初始化数据库，创建 chat_sessions 和 chat_messages 表。
会话管理：使用 save_chat_session 函数保存会话和消息。使用 load_all_sessions 和 load_session_messages 函数加载会话和消息。
导出功能：使用 export_session_to_csv 和 export_all_sessions_to_csv 函数导出会话为 CSV 文件。
清空数据库：使用 clear_database 函数清空数据库。
路由：
/: 主页路由，渲染 index_app06.html。
/chat: 处理聊天请求，保存会话和消息。
/load_sessions: 加载所有会话。
/load_session: 加载指定会话的消息。
/export_session_csv: 导出当前会话为 CSV。
/export_all_csv: 导出所有会话为 CSV。
/clear_db: 清空数据库。
"""

from flask import Flask, render_template, request, jsonify
import requests
import sqlite3
from datetime import datetime
import os
import csv

app = Flask(__name__)

# 全局变量存储当前数据库路径
dir_path = os.path.dirname(os.path.realpath(__file__))
current_db_path = os.path.join(dir_path, 'chat_history_app06.db')

# 初始化数据库
def init_db(db_path):
    conn = sqlite3.connect(db_path)
    c = conn.cursor()
    # 创建会话表
    c.execute('''CREATE TABLE IF NOT EXISTS chat_sessions
                 (id INTEGER PRIMARY KEY AUTOINCREMENT,
                  session_id TEXT UNIQUE,
                  topic TEXT,
                  created_at TEXT,
                  updated_at TEXT)''')
    # 创建消息表
    c.execute('''CREATE TABLE IF NOT EXISTS chat_messages
                 (id INTEGER PRIMARY KEY AUTOINCREMENT,
                  session_id TEXT,
                  role TEXT,
                  content TEXT,
                  timestamp TEXT,
                  FOREIGN KEY(session_id) REFERENCES chat_sessions(session_id))''')
    conn.commit()
    conn.close()

# 保存会话和消息
def save_chat_session(db_path, session_id, topic, messages):
    conn = sqlite3.connect(db_path)
    c = conn.cursor()
    # 插入会话
    created_at = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    c.execute("INSERT OR IGNORE INTO chat_sessions (session_id, topic, created_at, updated_at) VALUES (?, ?, ?, ?)",
              (session_id, topic, created_at, created_at))
    # 插入消息
    for message in messages:
        c.execute("INSERT INTO chat_messages (session_id, role, content, timestamp) VALUES (?, ?, ?, ?)",
                  (session_id, message['role'], message['content'], datetime.now().strftime('%Y-%m-%d %H:%M:%S')))
    conn.commit()
    conn.close()

# 加载所有会话
def load_all_sessions(db_path):
    conn = sqlite3.connect(db_path)
    c = conn.cursor()
    c.execute("SELECT session_id, topic, created_at FROM chat_sessions ORDER BY created_at DESC")
    sessions = c.fetchall()
    conn.close()
    return sessions

# 加载指定会话的消息
def load_session_messages(db_path, session_id):
    conn = sqlite3.connect(db_path)
    c = conn.cursor()
    c.execute("SELECT role, content, timestamp FROM chat_messages WHERE session_id = ? ORDER BY timestamp", (session_id,))
    messages = c.fetchall()
    conn.close()
    return messages

# 导出会话为 CSV
def export_session_to_csv(db_path, session_id, csv_path):
    messages = load_session_messages(db_path, session_id)
    with open(csv_path, mode='w', newline='', encoding='utf-8') as file:
        writer = csv.writer(file)
        writer.writerow(['Role', 'Content', 'Timestamp'])
        for message in messages:
            writer.writerow(message)

# 导出所有会话为 CSV
def export_all_sessions_to_csv(db_path, csv_path):
    sessions = load_all_sessions(db_path)
    with open(csv_path, mode='w', newline='', encoding='utf-8') as file:
        writer = csv.writer(file)
        writer.writerow(['Session ID', 'Topic', 'Created At'])
        for session in sessions:
            writer.writerow(session)

# 清空数据库
def clear_database(db_path):
    conn = sqlite3.connect(db_path)
    c = conn.cursor()
    c.execute("DELETE FROM chat_sessions")
    c.execute("DELETE FROM chat_messages")
    conn.commit()
    conn.close()

# 初始化数据库
init_db(current_db_path)

# 主页路由
@app.route('/')
def index():
    return render_template('index_app06.html')

# 聊天路由
@app.route('/chat', methods=['POST'])
def chat():
    global current_db_path

    user_input = request.json.get('message')
    session_id = request.json.get('session_id')
    session_messages = request.json.get('messages', [])

    # 如果是新会话，生成会话 ID 并设置主题
    if not session_id:
        session_id = datetime.now().strftime('%Y%m%d%H%M%S')
        topic = user_input  # 第一个用户输入作为主题
    else:
        topic = None  # 已有会话，不需要更新主题

    # 添加用户消息
    session_messages.append({"role": "user", "content": user_input})

    # 调用 API（这里假设使用本地模型）
    ai_response = "这是 AI 的回复"  # 替换为实际的 API 调用
    session_messages.append({"role": "assistant", "content": ai_response})

    # 保存会话和消息
    if topic:
        save_chat_session(current_db_path, session_id, topic, session_messages)
    else:
        save_chat_session(current_db_path, session_id, None, session_messages[-2:])  # 只保存最新两条消息

    return jsonify({"response": ai_response, "messages": session_messages, "session_id": session_id})

# 加载所有会话
@app.route('/load_sessions', methods=['GET'])
def load_sessions():
    sessions = load_all_sessions(current_db_path)
    return jsonify({"sessions": sessions})

# 加载指定会话的消息
@app.route('/load_session', methods=['GET'])
def load_session():
    session_id = request.args.get('session_id')
    messages = load_session_messages(current_db_path, session_id)
    return jsonify({"messages": messages})

# 导出当前会话为 CSV
@app.route('/export_session_csv', methods=['POST'])
def export_session_csv():
    session_id = request.json.get('session_id')
    csv_path = request.json.get('csv_path', f"chat_session_{session_id}.csv")
    export_session_to_csv(current_db_path, session_id, csv_path)
    return jsonify({"status": "success", "csv_path": csv_path})

# 导出所有会话为 CSV
@app.route('/export_all_csv', methods=['POST'])
def export_all_csv():
    csv_path = request.json.get('csv_path', "all_chat_sessions.csv")
    export_all_sessions_to_csv(current_db_path, csv_path)
    return jsonify({"status": "success", "csv_path": csv_path})

# 清空数据库
@app.route('/clear_db', methods=['POST'])
def clear_db():
    clear_database(current_db_path)
    return jsonify({"status": "success"})

if __name__ == '__main__':
    app.run(debug=True)